package manager

import (
	"sync"

	"github.com/gofiber/websocket/v2"
)

type WebSocketManager struct {
	Connections map[string]func(*websocket.Conn)
	mu          sync.RWMutex
	InitMessage map[string]string
	MutexMap    map[string]*sync.Mutex
}

func (m *WebSocketManager) Register(knotID string, handler func(*websocket.Conn)) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.MutexMap[knotID] = &sync.Mutex{}
	m.Connections[knotID] = handler
}

func (m *WebSocketManager) UnRegister(knotID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.Connections, knotID)
	delete(m.MutexMap, knotID)
	delete(m.InitMessage, knotID)
}

func (m *WebSocketManager) Handle(knotID string, conn *websocket.Conn) {
	m.mu.RLock()
	handler, exists := m.Connections[knotID]
	m.mu.RUnlock()

	if exists {
		handler(conn)
	} else {
		conn.WriteMessage(websocket.TextMessage, []byte("xxclosexx"))
		conn.Close()
	}
}

func (m *WebSocketManager) Remove(knotID string) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	_, exists := m.Connections[knotID]
	if exists {
		delete(m.Connections, knotID)
	}
}

func (m *WebSocketManager) HaveHandler(knotID string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	_, exists := m.Connections[knotID]
	return exists
}
